import os
import pandas as pd
from loguru import logger

def kucun_chengben_14_16_count_main(extract_dir):
    # 列出当前目录下的所有文件
    files_in_directory = os.listdir(extract_dir)
    logger.info(files_in_directory)

    if len(files_in_directory) == 1 and os.path.isdir(os.path.join(extract_dir, files_in_directory[0])):
        extract_dir = os.path.join(extract_dir, files_in_directory[0])
        files_in_directory = os.listdir(extract_dir)
        logger.info(files_in_directory)

    # 过滤出所有的 .xlsx 文件
    xlsx_file_name_list = [file for file in files_in_directory if not file.startswith(".~") and file.endswith('.xlsx')]
    logger.info(xlsx_file_name_list)
    logger.info("正在处理文件...")

    if xlsx_file_name_list and "HK_SMT_SA报表整理" in xlsx_file_name_list[0]:
        xlsx_file_name = xlsx_file_name_list[0]
        file_path = os.path.join(extract_dir, xlsx_file_name)
    else:
        raise FileNotFoundError("未找到包含 'HK_SMT_SA报表整理' 的 xlsx 文件")
    
    # 读取 Excel 文件
    xls = pd.ExcelFile(file_path)

    # 读取子表
    df_A = pd.read_excel(xls, 'bizWarehouseProductInventory')
    df_B = pd.read_excel(xls, 'Receiving-本月')
    df_C = pd.read_excel(xls, 'Orders-本月')
    df_D = pd.read_excel(xls, '上月check')

    #插入SKU列
    df_B['sku'] = df_B['产品代码'].str.split('-').str[-1]

    original_col_index = df_B.columns.get_loc('产品代码')
    new_col_index = original_col_index + 1
    df_B.insert(new_col_index, 'sku', df_B.pop('sku'))

    # 获取 sku 并集
    sku_union = pd.Series(list(set(df_A['sku']).union(set(df_D['sku']))))

    # 计算 上期
    df_check = pd.DataFrame()
    df_check['sku'] = sku_union
    df_check = df_check.merge(df_D[['sku', '下载数据期末']], on='sku', how='left').fillna(0)
    df_check.rename(columns={'下载数据期末': '上期'}, inplace=True)

    # 计算 接收入库
    df_check = df_check.merge(df_B.groupby('sku')['上架数量'].sum().reset_index(), on='sku', how='left').fillna(0)
    df_check.rename(columns={'上架数量': '接收入库'}, inplace=True)

    # 固定 退仓 为 0
    df_check['退仓'] = 0

    # 计算 销售出库
    df_C['数量1'] = df_C['数量1'].fillna(0)
    df_C['数量2'] = df_C['数量2'].fillna(0)
    sales_out = df_C.groupby('sku1')['数量1'].sum().reset_index().rename(columns={'sku1': 'sku', '数量1': '销售出库'})
    sales_out2 = df_C.groupby('sku2')['数量2'].sum().reset_index().rename(columns={'sku2': 'sku', '数量2': '销售出库'})
    sales_out = sales_out.set_index('sku').add(sales_out2.set_index('sku'), fill_value=0).reset_index()
    df_check = df_check.merge(sales_out, on='sku', how='left').fillna(0)

    # 计算 期末
    df_check['期末'] = df_check['上期'] + df_check['接收入库'] + df_check['退仓'] - df_check['销售出库']

    # 计算 下载数据期末
    df_check = df_check.merge(df_A.groupby('sku')['可售'].sum().reset_index(), on='sku', how='left').fillna(0)
    df_check.rename(columns={'可售': '下载数据期末'}, inplace=True)

    # 计算 DIF
    df_check['DIF'] = df_check['期末'] - df_check['下载数据期末']

    # 生成 拆分组合sku 子表
    split_sku_data = []

    for index, row in df_check.iterrows():
        sku = row['sku']
        if '-' in sku:
            parts = sku.split('-')
            if len(parts) == 2:
                split_sku_data.append({**row, 'sku-拆分后字段': parts[0]})
                split_sku_data.append({**row, 'sku-拆分后字段': parts[1]})
            elif len(parts) == 3 and parts[1].isdigit():
                for _ in range(int(parts[1])):
                    split_sku_data.append({**row, 'sku-拆分后字段': parts[0]})
                split_sku_data.append({**row, 'sku-拆分后字段': parts[2]})
        else:
            split_sku_data.append({**row, 'sku-拆分后字段': sku})

    df_split_sku = pd.DataFrame(split_sku_data)

    # 写入 Excel 文件
    with pd.ExcelWriter(file_path, mode='a') as writer:
        df_check.to_excel(writer, sheet_name='check', index=False)
        df_split_sku.to_excel(writer, sheet_name='拆分组合sku', index=False)


if __name__ == '__main__':
    kucun_chengben_14_16_count_main("./测试数据")
